#!/usr/bin/env python
import functools
import time
import py_trees


def pre_tick_handler(behaviour_tree: py_trees.trees.BehaviourTree) -> None:
    """Print a banner immediately before every tick of the tree.
    Args:
        behaviour_tree (:class:`~py_trees.trees.BehaviourTree`): the tree custodian
    """
    print("\n--------- Run %s ---------\n" % behaviour_tree.count)


def post_tick_handler(
    snapshot_visitor: py_trees.visitors.SnapshotVisitor,
    behaviour_tree: py_trees.trees.BehaviourTree,
) -> None:
    """Print an ascii tree with the current snapshot status."""
    print(
        py_trees.display.unicode_tree(
            root=behaviour_tree.root,
            visited=snapshot_visitor.visited,
            previously_visited=snapshot_visitor.previously_visited,
        )
    )


def create_root() -> py_trees.behaviour.Behaviour:
    """
    Create the root behaviour and it's subtree.
    Returns:
        the root behaviour
    """
    task_one = py_trees.behaviours.StatusQueue(
        name="Task 1",
        queue=[
            py_trees.common.Status.RUNNING,
            py_trees.common.Status.RUNNING,
        ],
        eventually=py_trees.common.Status.SUCCESS,
    )
    task_two = py_trees.behaviours.StatusQueue(
        name="Task 2",
        queue=[
            py_trees.common.Status.RUNNING,
            py_trees.common.Status.RUNNING,
        ],
        eventually=py_trees.common.Status.SUCCESS,
    )
    high_priority_interrupt = py_trees.decorators.RunningIsFailure(
        name="Running is Failure",
        child=py_trees.behaviours.Periodic(name="High Priority", n=3),
    )
    piwylo = py_trees.idioms.pick_up_where_you_left_off(name="Pick Up\nWhere You\nLeft Off", tasks=[task_one, task_two])
    root = py_trees.composites.Selector(name="Root", memory=False)
    root.add_children([high_priority_interrupt, piwylo])

    return root


def main() -> None:
    """Entry point for the demo script."""
    py_trees.logging.level = py_trees.logging.Level.DEBUG
    root = create_root()
    py_trees.display.render_dot_tree(root)

    ####################
    # Tree Stewardship
    ####################
    behaviour_tree = py_trees.trees.BehaviourTree(root)
    behaviour_tree.add_pre_tick_handler(pre_tick_handler)
    behaviour_tree.visitors.append(py_trees.visitors.DebugVisitor())
    snapshot_visitor = py_trees.visitors.SnapshotVisitor()
    behaviour_tree.add_post_tick_handler(functools.partial(post_tick_handler, snapshot_visitor))

    behaviour_tree.visitors.append(snapshot_visitor)
    behaviour_tree.setup(timeout=15)

    ####################
    # Tick Tock
    ####################
    interactive = False
    if interactive:
        py_trees.console.read_single_keypress()
    for _unused_i in range(1, 11):
        try:
            behaviour_tree.tick()
            if interactive:
                py_trees.console.read_single_keypress()
            else:
                time.sleep(0.5)
        except KeyboardInterrupt:
            break
    print("\n")

if __name__ == "__main__":
    main()
